#!/usr/bin/env python
"""
FRP 服务器实现
该模块实现了一个 FRP 服务器，用于处理客户端连接和端口转发。
主要功能：
1. 监听用户连接请求
2. 管理多个 FRP 客户端连接
3. 实现负载均衡
4. 处理端口转发
"""

import os
import subprocess
import sys
import socket
import time
import threading
import struct
import selectors
import logging
import traceback
from typing import List, Optional, Tuple, Dict
from collections import deque

import lib.ConnTool as ConnTool

# 配置日志
logger = logging.getLogger("FRP")

# 如果日志处理器已经存在，不重复添加
if not logger.handlers:
    formatter = logging.Formatter("%(asctime)s [%(levelname)s] %(message)s", "%H:%M:%S")
    console_handler = logging.StreamHandler(sys.stdout)
    console_handler.setFormatter(formatter)
    logger.addHandler(console_handler)
    logger.setLevel(logging.INFO)

# 创建全局选择器
selector = selectors.DefaultSelector()


class FrpClient:
    """
    FRP 客户端连接管理类

    属性:
        cmd_conn (socket.socket): 控制连接
        work_conns (List[socket.socket]): 工作连接池
        last_heartbeat (float): 最后心跳时间
        is_active (bool): 客户端是否活跃
    """

    def __init__(self, cmd_conn: socket.socket):
        self.cmd_conn = cmd_conn
        self.work_conns: List[socket.socket] = []
        self.last_heartbeat = time.time()
        self.is_active = True


class Frps(threading.Thread):
    """
    FRP 服务器类，继承自 Thread 以支持多线程处理

    属性:
        sock (socket.socket): 用户连接监听套接字
        frpc_sock (socket.socket): FRP 客户端连接监听套接字
        clients (Dict[socket.socket, FrpClient]): FRP 客户端连接池
        user_conns (List[socket.socket]): 用户连接池
        client_queue (deque): 客户端轮询队列
        client_ips (Dict[str, List[socket.socket]]): 按IP地址分组的客户端连接
    """

    def __init__(self, port: int, target_port: int):
        """
        初始化 FRP 服务器

        参数:
            port: 用户连接端口
            target_port: FRP 客户端连接端口
        """
        super().__init__()
        self.daemon = True  # 设置为守护线程

        # 初始化用户连接监听套接字
        self.sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        self.sock.bind(("0.0.0.0", port))
        self.sock.setblocking(False)
        self.sock.listen(100)
        selector.register(self.sock, selectors.EVENT_READ, self.accept_connection)

        # 初始化 FRP 客户端连接监听套接字
        self.frpc_sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        self.frpc_sock.bind(("0.0.0.0", target_port))
        self.frpc_sock.setblocking(False)
        self.frpc_sock.listen(100)
        selector.register(
            self.frpc_sock, selectors.EVENT_READ, self.accept_frp_connection
        )

        # 初始化连接池
        self.clients: Dict[socket.socket, FrpClient] = {}
        self.user_conns: List[socket.socket] = []
        self.client_queue = deque()  # 用于轮询的客户端队列
        self.client_ips: Dict[str, List[socket.socket]] = {}  # 按IP地址分组的客户端连接

        # 启动心跳检查线程
        threading.Thread(target=self.check_heartbeats, daemon=True).start()

    def check_heartbeats(self) -> None:
        """
        检查所有客户端的心跳状态
        如果超过30秒没有心跳，则认为客户端已断开
        """
        while True:
            try:
                current_time = time.time()
                for cmd_conn, client in list(self.clients.items()):
                    try:
                        client_addr = cmd_conn.getpeername()
                    except:
                        client_addr = "未知"
                    time_since_last_heartbeat = round(
                        current_time - client.last_heartbeat, 2
                    )
                    if time_since_last_heartbeat > 30:
                        try:
                            logger.info(
                                f"客户端 {client_addr} 心跳超时，最后心跳时间: {time_since_last_heartbeat}秒前"
                            )
                        except (socket.error, OSError):
                            logger.info("检测到已断开的客户端连接")
                        finally:
                            self.remove_client(cmd_conn)

                    logger.debug(
                        f"{client_addr} {client.is_active} {time.strftime('%Y-%m-%d %H:%M:%S', time.localtime(client.last_heartbeat))} {time_since_last_heartbeat}"
                    )

                logger.debug(f"当前可用客户端数: {self.get_client_count()}")
                time.sleep(9)  # 每9秒检查一次
            except Exception as e:
                logger.error(f"心跳检查错误: {e}")
                e_info = traceback.format_exc()
                logger.error(f"错误信息: {e_info}")
                time.sleep(8)

    def get_client_count(self) -> int:
        """
        获取当前活跃的客户端数量（按IP地址统计）

        返回:
            当前活跃的客户端数量
        """
        sum_client = len(self.client_ips)
        with open("client_ips.txt", "w") as f:
            f.write(str(sum_client))
        return sum_client

    def update_client_ips(self) -> None:
        """
        更新客户端IP分组，只保留存在于self.clients中的连接
        """
        try:
            # 遍历所有IP组
            for ip in list(self.client_ips.keys()):
                # 过滤掉不在clients中的连接
                valid_conns = [
                    conn for conn in self.client_ips[ip] if conn in self.clients
                ]
                if valid_conns:
                    self.client_ips[ip] = valid_conns
                else:
                    # 如果没有有效连接，删除该IP组
                    del self.client_ips[ip]
        except Exception as e:
            logger.error(f"更新客户端IP分组错误: {e}")

    def remove_client(self, cmd_conn: socket.socket) -> None:
        """
        移除客户端连接

        参数:
            cmd_conn: 要移除的客户端控制连接
        """
        try:
            # 先更新client_ips字典，确保只保留有效连接
            self.update_client_ips()

            if cmd_conn not in self.clients:  # 增加存在性检查
                return

            # 先尝试获取IP地址（在关闭连接前）
            client_ip = None
            try:
                if cmd_conn.fileno() != -1:
                    client_ip = cmd_conn.getpeername()[0]
            except Exception:
                pass

            # 增加连接状态检查
            if cmd_conn.fileno() != -1:  # <-- 新增有效性检查
                selector.unregister(cmd_conn)
                cmd_conn.close()

            # 从客户端池中移除
            del self.clients[cmd_conn]
            # 从轮询队列中移除
            if cmd_conn in self.client_queue:
                self.client_queue.remove(cmd_conn)

            # 从IP分组中移除（使用之前保存的IP地址）
            if client_ip and client_ip in self.client_ips:
                if cmd_conn in self.client_ips[client_ip]:
                    self.client_ips[client_ip].remove(cmd_conn)
                if not self.client_ips[client_ip]:
                    del self.client_ips[client_ip]

            # 再次更新client_ips确保全部同步
            self.update_client_ips()
        except Exception as e:
            logger.error(f"移除客户端错误: {e}")
            # 强制清理残留
            if cmd_conn in self.clients:
                del self.clients[cmd_conn]

    def get_next_client(self) -> Optional[FrpClient]:
        """
        获取下一个可用的客户端（轮询方式）

        返回:
            下一个可用的客户端，如果没有可用客户端则返回 None
        """
        if not self.client_queue:
            return None

        # 轮询直到找到活跃的客户端
        for _ in range(len(self.client_queue)):
            cmd_conn = self.client_queue[0]
            self.client_queue.rotate(-1)  # 将当前客户端移到队尾

            if cmd_conn in self.clients and self.clients[cmd_conn].is_active:
                return self.clients[cmd_conn]

        return None

    def heartbeat(self) -> None:
        """
        向所有活跃的 FRP 客户端发送心跳包
        每9秒发送一次心跳包以保持连接活跃
        """
        while True:
            try:
                if not self.is_connected:
                    logger.info("尝试重新连接服务器...")
                    try:
                        self.connect_to_server()
                    except:
                        self.reconnect_delay = min(
                            self.reconnect_delay * 2, self.max_reconnect_delay
                        )
                        logger.info(f"重连失败，{self.reconnect_delay}秒后重试")
                        time.sleep(self.reconnect_delay)
                        continue

                if self.server_fd is not None:
                    try:
                        self.server_fd.send(struct.pack("i", 1))
                        logger.debug("心跳包发送成功")
                    except Exception as e:
                        logger.error(f"心跳包发送错误: {e}")
                        self.is_connected = False
                        continue

                time.sleep(9)
            except Exception as e:
                logger.error(f"心跳包发送错误: {e}")
                time.sleep(1)

    def connect_to_server(self) -> None:
        try:
            if self.server_fd is not None:
                try:
                    selector.unregister(self.server_fd)
                    self.server_fd.close()
                except:
                    pass
                self.server_fd = None

            self.server_fd = socket.create_connection(
                (self.server_host, self.server_port)
            )
            self.server_fd.sendall(struct.pack("i", 1))
            self.server_fd.setblocking(False)
            selector.register(
                self.server_fd, selectors.EVENT_READ, self.handle_controller_data
            )
            self.is_connected = True
            self.reconnect_delay = 5
            logger.info("已连接到服务器")
        except Exception as e:
            self.is_connected = False
            logger.error(f"连接服务器失败: {e}")
            raise

    def accept_connection(self, sock: socket.socket, mask: int) -> None:
        """
        处理用户连接请求

        参数:
            sock: 监听套接字
            mask: 事件掩码
        """
        try:
            user_conn, addr = self.sock.accept()
            user_conn.setblocking(True)
            self.user_conns.append(user_conn)
            logger.info(f"收到用户请求，来自: {addr}")

            # 获取下一个可用的客户端
            client = self.get_next_client()
            if client is None:
                logger.warning("没有可用的 FRP 客户端，无法处理用户请求")
                return

            try:
                client.cmd_conn.send(struct.pack("i", 2))  # 发送建立新 TCP 连接的命令
            except IOError as err:
                logger.error(f"发送命令失败: {err}")
                client.is_active = False
        except Exception as e:
            logger.error(f"处理用户连接错误: {e}")

    def accept_frp_connection(self, sock: socket.socket, mask: int) -> None:
        """
        处理 FRP 客户端连接请求

        参数:
            sock: 监听套接字
            mask: 事件掩码
        """
        try:
            frpc_conn, addr = sock.accept()
            frpc_conn.setblocking(False)
            selector.register(
                frpc_conn, selectors.EVENT_READ, self.handle_controller_data
            )

            # 创建新的客户端对象
            client = FrpClient(frpc_conn)
            self.clients[frpc_conn] = client
            self.client_queue.append(frpc_conn)

            # 添加到IP分组
            client_ip = addr[0]
            if client_ip not in self.client_ips:
                self.client_ips[client_ip] = []
            self.client_ips[client_ip].append(frpc_conn)

            logger.info(
                f"FRP 客户端已连接: {addr}, 当前客户端数: {self.get_client_count()}"
            )
        except Exception as e:
            logger.error(f"处理 FRP 客户端连接错误: {e}")

    def handle_controller_data(self, frpc_conn: socket.socket, mask: int) -> None:
        """
        处理来自 FRP 客户端的控制数据

        参数:
            frpc_conn: FRP 客户端连接
            mask: 事件掩码
        """
        try:
            data = frpc_conn.recv(4)
            if data:
                cmd = struct.unpack("i", data)[0]
                logger.debug(f"收到命令: {cmd}")

                if cmd == 2:  # 建立工作 TCP 连接
                    if not self.user_conns:
                        logger.warning("没有待处理的用户连接")
                        return

                    selector.unregister(frpc_conn)  # 取消监听
                    user_conn = self.user_conns.pop()
                    frpc_conn.setblocking(True)

                    # 将工作连接添加到客户端的工作连接池
                    if frpc_conn in self.clients:
                        self.clients[frpc_conn].work_conns.append(frpc_conn)

                    ConnTool.join(user_conn, frpc_conn)
                    logger.info("已建立工作 TCP 连接")

                elif cmd == 1:  # 心跳包
                    if frpc_conn in self.clients:
                        self.clients[frpc_conn].last_heartbeat = time.time()
                        if not self.clients[frpc_conn].is_active:
                            self.clients[frpc_conn].is_active = True
                            logger.info(f"客户端 {frpc_conn.getpeername()} 重新活跃")

                elif cmd == 3:  # 断开连接命令
                    try:
                        client_addr = frpc_conn.getpeername()
                        logger.info(f"客户端 {client_addr} 断开连接")
                    except:
                        logger.info("未知客户端请求断开连接")

                    # 立即移除客户端
                    self.remove_client(frpc_conn)
                    logger.info(f"当前可用客户端数: {self.get_client_count()}")

        except IOError:
            # 非阻塞模式下，如果没有数据可读会抛出异常
            pass
        except Exception as e:
            logger.error(f"处理控制数据错误: {e}")
            if frpc_conn in self.clients:
                self.clients[frpc_conn].is_active = False

    def run(self) -> None:
        """主事件循环，处理所有连接事件"""
        try:
            while True:
                events = selector.select()
                for key, mask in events:
                    callback = key.data
                    callback(key.fileobj, mask)
        except KeyboardInterrupt:
            logger.info("服务器被用户中断")
        except Exception as e:
            logger.error(f"服务器运行错误: {e}")
        finally:
            # 清理所有连接
            try:
                # 关闭所有客户端连接
                for cmd_conn in list(self.clients.keys()):
                    self.remove_client(cmd_conn)

                # 关闭监听套接字
                try:
                    selector.unregister(self.sock)
                    self.sock.close()
                except:
                    pass

                try:
                    selector.unregister(self.frpc_sock)
                    self.frpc_sock.close()
                except:
                    pass

                # 关闭选择器
                selector.close()
            except:
                pass

            logger.info("服务器已停止")


def kill_process_by_port(port):
    try:
        if os.name == "posix":  # Linux/macOS
            # 使用 lsof 命令，它更可靠
            cmd = f"lsof -i :{port} -t"
            result = subprocess.run(
                cmd,
                shell=True,
                stdout=subprocess.PIPE,
                stderr=subprocess.PIPE,
                text=True,
            )
            if result.stdout.strip():
                pids = result.stdout.strip().split("\n")
                for pid in pids:
                    if pid.isdigit():
                        logger.info(f"发现端口 {port} 被进程 {pid} 占用，正在终止...")
                        try:
                            os.kill(int(pid), 9)
                        except ProcessLookupError:
                            pass  # 进程可能已经不存在

                # 等待端口释放，最多等待10秒
                for _ in range(10):
                    if is_port_available(port):
                        logger.info(f"端口 {port} 已成功释放")
                        return
                    time.sleep(1)
                logger.warning(f"端口 {port} 可能未完全释放")
        else:  # Windows
            cmd = f"netstat -ano | findstr :{port}"
            result = subprocess.run(
                cmd,
                shell=True,
                stdout=subprocess.PIPE,
                stderr=subprocess.PIPE,
                text=True,
            )
            if result.stdout.strip():
                pid = result.stdout.split()[-1]  # 提取 PID
                subprocess.run(f"taskkill /PID {pid} /F", shell=True)
                time.sleep(2)
    except Exception as e:
        logger.warning(f"处理端口 {port} 时出错: {e}")


# 检查端口是否真的可用
def is_port_available(port):
    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
        return s.connect_ex(("0.0.0.0", port)) != 0


def main() -> None:
    """主函数入口"""
    try:
        with open("client_ips.txt", "w") as f:
            f.write(str(0))

        # ============= 配置参数 =============
        # 服务器端口配置
        frps_port = xxxx  # FRP服务器监听端口
        user_port = xxxx  # 用户连接端口
        # ===================================

        # 检查并释放端口
        for port in [frps_port, user_port]:
            kill_process_by_port(port)
            if not is_port_available(port):
                logger.error(f"端口 {port} 仍被占用，请手动关闭冲突程序！")

        # 确保端口已经释放
        time.sleep(2)

        logger.info("正在启动 FRP 服务器...")
        server = Frps(user_port, frps_port)
        server.start()
        logger.info(
            f"FRP 服务器已启动 - 监听地址: 0.0.0.0:{frps_port}, 用户端口: {user_port}"
        )

        # 保持主线程运行
        server.join()

    except Exception as e:
        e_info = traceback.format_exc()
        logger.error(f"启动失败: {e_info}")


if __name__ == "__main__":
    main()
